from app.models import UserDB
import os
import argparse

def init_database(admin_username='admin', admin_password='admin123', force=False):
    """初始化数据库"""
    try:
        # 创建数据库实例
        db = UserDB()
        
        # 检查是否已存在管理员账户
        conn = db._get_db_connection()
        c = conn.cursor()
        c.execute('SELECT * FROM users WHERE username = ? AND is_admin = 1', (admin_username,))
        admin_exists = c.fetchone()
        
        if not admin_exists or force:
            if force and admin_exists:
                # 如果强制重新创建，先删除现有管理员账户
                c.execute('DELETE FROM users WHERE username = ? AND is_admin = 1', (admin_username,))
                conn.commit()
            
            # 创建管理员账户
            if db.create_user(admin_username, admin_password, is_admin=1):
                print(f"成功创建管理员账户:")
                print(f"用户名: {admin_username}")
                print(f"密码: {'*' * len(admin_password)}")
            else:
                print("创建管理员账户失败")
        else:
            print("管理员账户已存在，跳过创建")
            
        conn.close()
        print("数据库初始化完成！")
        
    except Exception as e:
        print(f"数据库初始化失败: {str(e)}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='初始化文件传输工具数据库')
    parser.add_argument('--force', action='store_true', help='强制重新创建管理员账户')
    parser.add_argument('--username', default='admin', help='管理员用户名')
    parser.add_argument('--password', default='admin123', help='管理员密码')
    
    args = parser.parse_args()
    init_database(args.username, args.password, args.force)
